package com.linovi.tree

class Tree(val data: Int, val left: Option[Tree] = None, val right: Option[Tree] = None) {

  def +(data: Int): Tree = {
    def insert(optTree: Option[Tree]): Tree = optTree match {
      case None => Tree(data)
      case Some(tree) => if (data < tree.data) tree(left = Some(insert(tree.left)))
      else if (data > tree.data) tree(right = Some(insert(tree.right)))
      else tree
    }
    insert(Some(this))
  }

  def search(data: Int) = {
    def search(optTree: Option[Tree]): Option[Tree] = optTree match {
      case None => None
      case Some(tree) => if (data < tree.data) search(tree.left)
      else if (data > tree.data) search(tree.right)
      else optTree
    }
    search(Some(this))
  }

  def contains(data: Int) = search(data) != None

  def -(data: Int): Tree = {
    def remove(optTree: Option[Tree]): Tree = optTree match {
      case None => this
      case Some(tree) => if (data < tree.data) tree(left = Some(remove(tree.left)))
      else if (data > tree.data) tree(right = Some(remove(tree.right)))
      else {
        (tree.hasLeft, tree.hasRight) match {
          case (false, false) => tree
          case (true, false) => tree.left.get
          case (false, true) => tree.right.get
          case (true, true) =>
            val rightMinData = tree.rightMin.get.data
            tree(data = rightMinData, right = Some(tree.right.get - rightMinData))
        }
      }
    }
    remove(Some(this))
  }

  /* def -(data: Int): Option[Tree] = {
    def remove(optTree: Option[Tree]): Option[Tree] = optTree match {
      case None => None
      case Some(tree) => if (data < tree.data) Some(tree(left = remove(tree.left)))
      else if (data > tree.data) Some(tree(right = remove(tree.right)))
      else {
        (tree.hasLeft, tree.hasRight) match {
          case (false, false) => None
          case (true, false) => tree.left
          case (false, true) => tree.right
          case (true, true) =>
            val rightMinData = tree.rightMin.get.data
            Some(tree(data = rightMinData, right = tree.right.get - rightMinData))
        }
      }
    }
    remove(Option(this))
  }*/

  def rightMin = {
    if (hasRight) {
      def rightMin(optTree: Option[Tree]): Option[Tree] =
        if (!optTree.get.hasLeft) optTree else rightMin(optTree.get.left)
      rightMin(right)
    } else None
  }

  val (isLeaf, hasLeft, hasRight) = (left, right) match {
    case (None, None) => (true, false, false)
    case (Some(_), None) => (false, true, false)
    case (None, Some(_)) => (false, false, true)
    case (Some(_), Some(_)) => (false, true, true)
  }

  def filter(f: (Int) => Boolean): Tree = {
    def filter(t: Tree, f: (Int) => Boolean): Tree = {
      if (!f(t.data)) {
        val newTree = t - t.data
        (newTree.hasLeft, newTree.hasRight) match {
          case (false, false) => newTree
          case (true, false) => filter(newTree.left.get, f)
          case (false, true) => filter(newTree.right.get, f)
          case (true, true) => newTree
        }
      } else t
    }
    filter(this, f)
  }

  def map(f: (Tree) => Tree): Tree = {
    def map(r: Tree, t: Option[Tree], f: (Tree) => Tree): Tree = {
      t match {
        case None => r
        case Some(tree) => map(map(r + f(tree).data, tree.right, f), tree.left, f)
      }
    }
    map(new Tree(data), Option(this), f)
  }

  def apply(data: Int = data, left: Option[Tree] = left, right: Option[Tree] = right) = {
    new Tree(data = data, left = left, right = right)
  }

  override def toString = {
    def childToString(child: Option[Tree], label: String) = child match {
      case None => ""
      case Some(tree) => label + tree.toString
    }
    val children =
      if (isLeaf) data.toString
      else data.toString + "(" + childToString(left, "l = ") + (if (hasLeft && hasRight) ", " else "") + childToString(right, "r = ") + ")"
    children
  }
}

object Tree {
  def apply(root: Int, data: Int*) = {
    def loop(tree: Tree, i: Int): Tree =
      if (i > data.length - 1) tree
      else loop(tree + data(i), i + 1)
    loop(new Tree(data = root), 0)
  }
}